# test_app.py - 简单的Flask应用用于测试花艺AI系统

from flask import Flask, render_template, request, jsonify, redirect, url_for
import requests
import os
import base64
from werkzeug.utils import secure_filename

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max upload

# 确保上传文件夹存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# 花艺AI API的URL
API_URL = "http://localhost:8000"  # 修改为你的API实际地址

@app.route('/')
def index():
    """主页，显示查询界面"""
    return render_template('index.html')

@app.route('/query', methods=['POST'])
def query():
    """处理查询请求"""
    query_text = request.form.get('query', '')
    generate_image = 'generate_image' in request.form
    
    # 调用API
    response = requests.post(
        f"{API_URL}/query",
        json={"query": query_text, "generate_image": generate_image}
    )
    
    result = response.json()
    return render_template('result.html', query=query_text, result=result)

@app.route('/upload', methods=['GET', 'POST'])
def upload():
    """处理知识库文件上传"""
    if request.method == 'POST':
        if 'file' not in request.files:
            return jsonify({"error": "No file part"})
            
        file = request.files['file']
        if file.filename == '':
            return jsonify({"error": "No selected file"})
            
        if file:
            filename = secure_filename(file.filename)
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)
            
            # 调用API上传文件
            with open(filepath, 'rb') as f:
                files = {"file": (filename, f)}
                response = requests.post(f"{API_URL}/upload-knowledge", files=files)
            
            # 删除临时文件
            os.remove(filepath)
            
            return jsonify(response.json())
    
    return render_template('upload.html')

@app.route('/status')
def status():
    """查看系统状态"""
    response = requests.get(f"{API_URL}/system-status")
    return render_template('status.html', status=response.json())

@app.route('/train')
def train():
    """触发本地模型训练"""
    response = requests.post(f"{API_URL}/train-local-model")
    return jsonify(response.json())

@app.route('/image', methods=['GET', 'POST'])
def generate_image():
    """生成花艺图像界面"""
    if request.method == 'POST':
        prompt = request.form.get('prompt', '')
        style = request.form.get('style', 'realistic')
        
        # 调用API生成图像
        response = requests.post(
            f"{API_URL}/generate-image",
            json={"prompt": prompt, "style": style}
        )
        
        result = response.json()
        return render_template('image_result.html', result=result)
    
    return render_template('generate_image.html')

if __name__ == '__main__':
    app.run(debug=True, port=5000)
